#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/5/21 14:43
# @File    : use_estimator.py.py
# @Author  : lxy
from pyspark.ml.feature import PolynomialExpansion, VectorAssembler
from pyspark.ml.linalg import Vectors
from pyspark.sql import SparkSession
from pyspark.ml.regression import LinearRegression

spark = SparkSession.builder.appName("Userestimator").getOrCreate()
spark.sparkContext.setLogLevel("WARN")
df = spark.read.csv("/master/train.csv",encoding="utf8",header=True,inferSchema=True,sep=",")
print("----------转换前-----------")
df.show()
df.printSchema()
vec_ass = VectorAssembler(inputCols=['x'],outputCol="feature")
feature_df = vec_ass.transform(df)
px = PolynomialExpansion(degree=2).setInputCol('feature').setOutputCol('polyfeature')
feature_df.show()
feature_df.printSchema()
feature_df = px.transform(feature_df)
print("----------转换后-----------")
feature_df.show()
feature_df.printSchema()
model_df=feature_df.select('polyfeature','y')
train_df,test_df=model_df.randomSplit([0.7,0.3])
print((train_df.count(), len(train_df.columns)))
print((test_df.count(), len(test_df.columns)))
print('-------------- 构建多项式回归模型 ------------------')
train_df.printSchema()
lin_Reg=LinearRegression(featuresCol='polyfeature',labelCol='y',predictionCol='pred',maxIter=200)
lr_model=lin_Reg.fit(train_df)
print('{}{}'.format('方程截距:',lr_model.intercept))
print('{}{}'.format('方程参数系数:',lr_model.coefficients))
training_predictions=lr_model.evaluate(train_df)
print('{}{}'.format('误差差值平方:',training_predictions.meanSquaredError))
print('{}{}'.format('判定系数：',training_predictions.r2 ))
test_results=lr_model.evaluate(test_df)
print(test_results.r2)
print(test_results.meanSquaredError)
